# coding=utf-8
##
## Author: jmdvirus@aliyun.com
##
## Create: 2019年02月14日 星期四 16时26分47秒
##

import numpy as np
import logging
import mxnet as mx

logging.getLogger().setLevel(logging.DEBUG)

batch_size = 50

def cba(src, suffix, num_filter, kernel, pad, stride=(1,1)):
    conv = mx.sym.Convolution(src, name="conv"+suffix, kernel=(kernel, kernel),
            pad=(pad, pad), num_filter=num_filter, stride=stride)
    bn = mx.sym.BatchNorm(conv, name="bn" +suffix, fix_gamma = False)
    act = mx.sym.Activation(bn, name="act" + suffix, act_type="relu")
    return act

net = mx.symbol.Variable('data')
net = cba(net, "1", 96, 3, 1)
#net = cba(net, "2", 96, 3, 1)
#net = cba(net, "3", 96, 3, 1, (2,2))
#net = cba(net, "4", 192, 3, 1)
#net = cba(net, "5", 192, 3, 1)
#net = cba(net, "6", 192, 3, 1, (2, 2))
#net = cba(net, "7", 192, 3, 0)
#net = cba(net, "8", 192, 1, 0)
#net = cba(net, "9", 10, 1, 0)

net = mx.sym.Pooling(net, name="pool", global_pool=True, 
        pool_type="avg", kernel=(1,1))

net = mx.sym.Flatten(net, name="flatten")
net = mx.sym.SoftmaxOutput(net, name="softmax")

shape = {"data": (batch_size, 3, 28, 28) }
mx.viz.print_summary(symbol=net, shape=shape)

def shownet(net):
    mx.viz.plot_network(symbol=net, shape=shape).view()

module = mx.mod.Module(symbol=net, context=mx.cpu(0))

prefix="/opt/data/data/ai.cifar/cifar"

train_iter = mx.io.ImageRecordIter(
        path_imgrec = prefix+"/otrain/train.rec",
        data_shape = (3, 28, 28),
        batch_size = batch_size,
        shuffle = True,
        rand_crop = True,
        rand_mirror = True,
        random_h = 10,
        random_s = 20,
        random_l = 25,
        max_random_scale = 1.20,
        min_random_scale = 0.88,
        max_rotate_angle = 20,
        max_aspect_ratio = 0.15,
        max_shear_ratio = 0.10,
        fill_value = 0,
        )

val_iter = mx.io.ImageRecordIter(
        path_imgrec = prefix+"/otest/test.rec",
        data_shape = (3, 28, 28),
        batch_size = batch_size,
        shuffle = False,
        rand_crop = False,
        rand_mirror = False,
        )

def totrain():
    module.fit(
            train_iter,
            eval_data = val_iter,
            initializer = mx.init.MSRAPrelu(slope=0.0),
            optimizer = 'sgd',
            optimizer_params = {'learning_rate':0.5,
                'lr_scheduler':mx.lr_scheduler.FactorScheduler(step=50000/batch_size,
                    factor=0.98)},
            num_epoch = 200,
            batch_end_callback = mx.callback.Speedometer(batch_size, 50000/batch_size)
            )

if __name__ == "__main__":
    #shownet(net)
    totrain()

